#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
数据库管理模块
用于连接MySQL数据库并读取训练数据
"""

import json
import mysql.connector
import os
from mysql.connector import Error
from typing import List, Dict, Any, Optional
import logging

# 配置日志
# 确保日志目录存在，使用os.makedirs的exist_ok参数避免竞争条件
log_dir = "logs"
os.makedirs(log_dir, exist_ok=True)

# 为这个模块创建独立的logger
logger = logging.getLogger('database_manager')
logger.setLevel(logging.INFO)

# 清除可能存在的旧处理器
for handler in logger.handlers[:]:
    logger.removeHandler(handler)

# 创建文件处理器
file_handler = logging.FileHandler(os.path.join(log_dir, "database_manager.log"), encoding='utf-8')
file_handler.setLevel(logging.INFO)

# 创建控制台处理器
console_handler = logging.StreamHandler()
# 设置控制台日志级别为INFO或以上
console_handler.setLevel(logging.INFO)

# 创建格式化器，增加进程和线程信息以便调试
formatter = logging.Formatter(
    '%(asctime)s - %(name)s - %(levelname)s - [PID: %(process)d, Thread: %(thread)d] - %(message)s'
)

file_handler.setFormatter(formatter)
console_handler.setFormatter(formatter)

# 添加处理器到logger
logger.addHandler(file_handler)
logger.addHandler(console_handler)

# 防止日志传播到根logger
logger.propagate = False

class DatabaseManager:
    """
    数据库管理器，用于连接MySQL数据库并读取训练数据
    """
    
    def __init__(self, host: str = "117.72.102.95", port: int = 16033, 
                 database: str = "yuyan", user: str = "yuyan", password: str = "yuyan@p3"):
        """
        初始化数据库连接参数
        
        Args:
            host: 数据库主机地址
            port: 数据库端口
            database: 数据库名称
            user: 用户名
            password: 密码
        """
        self.host = host
        self.port = port
        self.database = database
        self.user = user
        self.password = password
        self.connection = None
    
    def connect(self) -> bool:
        """
        建立数据库连接
        
        Returns:
            连接是否成功
        """
        try:
            logger.info(f"尝试连接数据库: {self.host}:{self.port}/{self.database}")
            self.connection = mysql.connector.connect(
                host=self.host,
                port=self.port,
                database=self.database,
                user=self.user,
                password=self.password,
                charset='utf8mb4',
                collation='utf8mb4_bin'
            )
            
            if self.connection.is_connected():
                logger.info("成功连接到MySQL数据库")
                db_info = self.connection.get_server_info()
                logger.info(f"MySQL服务器版本: {db_info}")
                return True
                
        except Error as e:
            logger.error(f"连接MySQL数据库时出错: {e}", exc_info=True)
            return False
    
    def disconnect(self):
        """
        关闭数据库连接
        """
        if self.connection and self.connection.is_connected():
            self.connection.close()
            logger.info("MySQL数据库连接已关闭")
    
    def get_bert_train_data(self, limit: Optional[int] = None) -> List[Dict[str, Any]]:
        """
        从bert_train表中获取训练数据
        
        Args:
            limit: 限制返回的记录数
            
        Returns:
            训练数据列表
        """
        if not self.connection or not self.connection.is_connected():
            if not self.connect():
                return []
        
        try:
            cursor = self.connection.cursor(dictionary=True)
            
            query = "SELECT id, field_id, original_text, entity_info FROM bert_train WHERE del_flag = 'N'"
            if limit:
                query += f" LIMIT {limit}"
            
            logger.info(f"执行查询: {query}")
            cursor.execute(query)
            records = cursor.fetchall()
            
            logger.info(f"查询返回 {len(records)} 条记录")
            
            # 处理entity_info字段，确保它是有效的JSON对象或数组
            processed_count = 0
            error_count = 0
            
            for i, record in enumerate(records):
                if isinstance(record['entity_info'], str):
                    try:
                        # 尝试解析JSON字符串
                        parsed_info = json.loads(record['entity_info'])
                        record['entity_info'] = parsed_info
                        processed_count += 1
                        logger.debug(f"记录 {record['id']} 的entity_info已解析为JSON对象")
                    except json.JSONDecodeError as e:
                        logger.warning(f"记录 {record['id']} 的entity_info不是有效的JSON格式: {record['entity_info'][:100]}... 错误: {e}")
                        record['entity_info'] = {}  # 设置为空字典作为默认值
                        error_count += 1
                elif record['entity_info'] is None:
                    logger.debug(f"记录 {record['id']} 的entity_info为None，设置为空字典")
                    record['entity_info'] = {}
                    processed_count += 1
                elif not isinstance(record['entity_info'], dict):
                    logger.warning(f"记录 {record['id']} 的entity_info是意外的类型: {type(record['entity_info'])}")
                    record['entity_info'] = {}
                    processed_count += 1
                else:
                    processed_count += 1
            
            logger.info(f"entity_info处理完成: {processed_count} 条记录处理成功, {error_count} 条记录处理失败")
            logger.info(f"成功获取 {len(records)} 条训练数据")
            return records
            
        except Error as e:
            logger.error(f"查询训练数据时出错: {e}", exc_info=True)
            return []
        except Exception as e:
            logger.error(f"处理训练数据时发生未预期的错误: {e}", exc_info=True)
            return []
        finally:
            if cursor:
                cursor.close()
    
    def get_bert_fields(self) -> List[Dict[str, Any]]:
        """
        获取领域信息
        
        Returns:
            领域信息列表
        """
        if not self.connection or not self.connection.is_connected():
            if not self.connect():
                return []
        
        try:
            cursor = self.connection.cursor(dictionary=True)
            cursor.execute("SELECT * FROM bert_field WHERE del_flag = 'N'")
            records = cursor.fetchall()
            
            logger.info(f"成功获取 {len(records)} 条领域信息")
            return records
            
        except Error as e:
            logger.error(f"查询领域信息时出错: {e}", exc_info=True)
            return []
        finally:
            if cursor:
                cursor.close()
    
    def get_bert_field_details(self, field_id: Optional[int] = None) -> List[Dict[str, Any]]:
        """
        获取领域详情信息
        
        Args:
            field_id: 领域ID，如果提供则只返回该领域的详情
            
        Returns:
            领域详情信息列表
        """
        if not self.connection or not self.connection.is_connected():
            if not self.connect():
                return []
        
        try:
            cursor = self.connection.cursor(dictionary=True)
            
            query = "SELECT * FROM bert_field_detail WHERE del_flag = 'N'"
            if field_id:
                query += f" AND field_id = {field_id}"
            
            logger.debug(f"执行查询: {query}")
            cursor.execute(query)
            records = cursor.fetchall()
            
            logger.info(f"成功获取 {len(records)} 条领域详情信息")
            return records
            
        except Error as e:
            logger.error(f"查询领域详情信息时出错: {e}", exc_info=True)
            return []
        finally:
            if cursor:
                cursor.close()
    
    def get_bert_nouns(self, field_id: Optional[int] = None) -> List[Dict[str, Any]]:
        """
        获取领域名词信息
        
        Args:
            field_id: 领域ID，如果提供则只返回该领域的名词
            
        Returns:
            领域名词信息列表
        """
        if not self.connection or not self.connection.is_connected():
            if not self.connect():
                return []
        
        try:
            cursor = self.connection.cursor(dictionary=True)
            
            query = "SELECT * FROM bert_noun WHERE del_flag = 'N'"
            if field_id:
                query += f" AND field_id = {field_id}"
            
            logger.debug(f"执行查询: {query}")
            cursor.execute(query)
            records = cursor.fetchall()
            
            logger.info(f"成功获取 {len(records)} 条领域名词信息")
            return records
            
        except Error as e:
            logger.error(f"查询领域名词信息时出错: {e}", exc_info=True)
            return []
        finally:
            if cursor:
                cursor.close()
    
    def get_field_code_by_id(self, field_id: int) -> Optional[str]:
        """
        根据领域ID获取领域代码
        
        Args:
            field_id: 领域ID
            
        Returns:
            领域代码
        """
        if not self.connection or not self.connection.is_connected():
            if not self.connect():
                return None
        
        try:
            cursor = self.connection.cursor(dictionary=True)
            logger.debug(f"查询领域ID {field_id} 的代码")
            cursor.execute("SELECT code FROM bert_field WHERE id = %s AND del_flag = 'N'", (field_id,))
            record = cursor.fetchone()
            
            if record:
                logger.debug(f"领域ID {field_id} 的代码为: {record['code']}")
                return record['code']
            else:
                logger.warning(f"未找到ID为 {field_id} 的领域信息")
                return None
                
        except Error as e:
            logger.error(f"查询领域代码时出错: {e}", exc_info=True)
            return None
        finally:
            if cursor:
                cursor.close()

# 测试代码
if __name__ == "__main__":
    # 创建数据库管理器实例
    db_manager = DatabaseManager()
    
    # 连接数据库
    if db_manager.connect():
        # 获取训练数据示例
        train_data = db_manager.get_bert_train_data(5)  # 限制获取5条记录
        print("训练数据示例:")
        for data in train_data:
            print(f"  ID: {data['id']}, 原文: {data['original_text'][:50]}...")
            if data['entity_info']:
                print(f"    实体信息: {data['entity_info']}")
            else:
                print("    实体信息: 无")
        
        # 获取领域信息示例
        fields = db_manager.get_bert_fields()
        print(f"\n领域信息数量: {len(fields)}")
        
        # 获取领域详情信息示例
        field_details = db_manager.get_bert_field_details()
        print(f"领域详情信息数量: {len(field_details)}")
        
        # 获取领域名词信息示例
        nouns = db_manager.get_bert_nouns()
        print(f"领域名词信息数量: {len(nouns)}")
        
        # 断开连接
        db_manager.disconnect()
    else:
        print("无法连接到数据库")